package com.yww.datastructure.btree;

import java.util.Arrays;

/**
 * @author yww
 * @description B 树
 * @since 2024/3/6 15:12
 */
public class BTree {
    Node root; // 根节点
    int t; // 树中节点最小的度（最小为2）
    final int MIN_KEY_NUMBER; // 树节点中最少key数目
    final int MAX_KEY_NUMBER; // 树节点中最大key数目

    public BTree() {
        this(2);
    }

    public BTree(int t) {
        this.t = t;
        root = new Node(t);
        MIN_KEY_NUMBER = t - 1;
        MAX_KEY_NUMBER = 2 * t - 1;

    }

    // 判断是否存在key
    public boolean contains(int key) {
        return root.get(key) != null;
    }

    public void put(int key) {
        doPut(root, key, null, 0);
    }

    // node：插入节点，key：插入关键字，parent：父节点，index：node的索引
    private void doPut(Node node, int key, Node parent, int index) {
        int i = 0;
        while (i < node.keyNumber) {
            if (node.keys[i] == key) { // 找到重复key，更新
                return;
            }
            if (node.keys[i] > key) { // 找到插入位置
                break;
            }
            i++;
        }
        if (node.leaf) { // 如果是叶子节点，则直接可以进行插入
            node.insertKey(i, key);
        } else { // 否则，继续在children[i]处继续递归插入
            doPut(node.children[i], key, node, i);
        }

        // 检查node节点是否要进行分裂（叶子节点插入、非叶子节点递归返回都要判断）
        if (node.keyNumber >= MAX_KEY_NUMBER) { // 达到上限进行分裂
            split(node, parent, index);
        }
    }

    // 分裂节点：left，父节点：parent，left的索引：index
    void split(Node left, Node parent, int index) {
        // 0. 要分裂节点为根节点
        if (root == left) {
            // 创建一个新根
            Node newRoot = new Node(t);
            newRoot.leaf = false;
            newRoot.insertChild(0, left);
            root = newRoot;
            parent = newRoot;
        }

        // 1. 创建 right 节点（分裂后大于当前 left 节点的），把 t 以后的 key 和 child 都拷贝过去
        Node right = new Node(t);
        // right和left为同一层，所以leaf一样
        right.leaf = left.leaf;
        // 将left的t以后的key转移到right
        System.arraycopy(left.keys, t, right.keys, 0, t - 1);
        if (!left.leaf) {
            // 分裂节点为非叶子，要把t以后的孩子也要移动到right，孩子比key多移动一个
            System.arraycopy(left.children, t, right.children, 0, t);
        }
        // right有效key数为转移个数t-1
        right.keyNumber = t - 1;
        // left由于要把key转移到parent和right，有效key数为t-1
        left.keyNumber = t - 1;

        // 2. t-1 处的 key 插入到 parent 的 index 处，index 指 left 作为孩子时的索引
        int mid = left.keys[t - 1];
        parent.insertKey(index, mid);

        // 3. right 节点作为 parent 的孩子插入到 index + 1 处
        parent.insertChild(index + 1, right);
    }

    // 删除方法
    public void remove(int key) {
        doRemove(null, root, 0, key);

    }

    // parnt：起始节点的父节点，node：从node开始进行删除，index：node的索引，key：要删除的key
    private void doRemove(Node parnt, Node node, int index, int key) {
        int i = 0;
        while (i < node.keyNumber) {
            if (node.keys[i] >= key) {
                break;
            }
            i++;
        }
        // i：找到，删除key在node的索引
        // i：没找到，往第i个child继续删除操作
        if (node.leaf) { // 叶子节点
            if (!(i < node.keyNumber && node.keys[i] == key)) { // case1：叶子节点没找到
                // 什么也不做，也不能进行平衡调整
                return;
            } else { // case2：叶子节点找到
                // 节点正常删除
                node.removeKey(i);
            }
        } else { // 非叶子
            if (!(i < node.keyNumber && node.keys[i] == key)) { // case3：非叶子节点没找到
                // 往第i个孩子继续递归删除
                doRemove(node, node.children[i], i, key);
            } else { // case4：非叶子节点找到
                // 寻找后继节点
                Node s = node.children[i + 1];
                while (!s.leaf) {
                    s = s.children[0];
                }
                // 找到后继key
                int sKey = s.keys[0];
                // 替换后继key
                node.keys[i] = sKey;
                // 删除后继key
                doRemove(node, node.children[i + 1], i + 1, sKey);
            }
        }

        if (node.keyNumber < MIN_KEY_NUMBER) { // 节点的key数目小于下限，进行平衡调整
            // case5、case6平衡调整
            balance(parnt, node, index);
        }
    }

    // 平衡调整，parent为x的父节点，x为调整节点，i为x的索引
    private void balance(Node parent, Node x, int i) {
        if (x == root) { // case6：调整节点为root节点，特殊处理
            if (root.keyNumber == 0 && root.children[0] != null) {
                root = root.children[0];
            }
            return;
        }
        Node left = parent.childLeftSibling(i);
        Node right = parent.childRightSibling(i);
        if (left != null && left.keyNumber > MIN_KEY_NUMBER) { // case5-1 左兄弟富裕，右旋借
            // 父亲i-1处的key右旋下，到x的0号位置
            x.insertKey(0, parent.keys[i - 1]);
            // 左兄弟最右侧的key，移除，右旋上到父节点i-1处
            parent.keys[i - 1] = left.removeRightmostKey();
            // 左兄弟不是叶子节点，将最右侧的孩子移出，到被调整节点0号位置
            if (!left.leaf) {
                x.insertChild(0, left.removeRightmostChild());
            }
            return;
        }
        if (right != null && right.keyNumber > MIN_KEY_NUMBER) { // case5-2 右兄弟富裕，左旋借
            // 父节点中后继key左旋转下来
            x.insertKey(parent.keys[i], x.keyNumber);
            // right中最小的key左旋上去
            parent.keys[i] = right.removeLeftmostKey();
            // right中最小的孩子换爹
            if (!right.leaf) {
                x.insertChild(x.keyNumber + 1, right.removeLeftmostChild());
            }
            return;
        }
        // case5-3：两边都不富裕，向左合并
        if (left != null) { // 有左兄弟，向左兄弟合并
            // 父节点移出被调整节点
            parent.removeChild(i);
            // 左兄弟插入父节点的i-1的key（x的前驱）
            left.insertKey(left.keyNumber, parent.removeKey(i - 1));
            // 转移x的所有key和child到左兄弟
            x.moveToTarget(left);
        } else { // 没有左兄弟，向自己合并
            // 父节点移出 i + 1 节点（被调整节点的右兄弟）
            parent.removeChild(i + 1);
            // 父节点的 i 处的key，合并到自己
            x.insertKey(x.keyNumber, parent.removeKey(i));
            // right所有key和child，合并到自己
            right.moveToTarget(x);
        }
    }

    static class Node {
        int[] keys; // 关键字
        Node[] children; // 孩子
        int keyNumber; // 有效关键字数目
        boolean leaf = true; // 是否是叶子节点
        int t; // 最小度数（最小孩子数）

        public Node(int t) {
            if (t < 2) t = 2; // t >= 2
            this.t = t;
            this.children = new Node[t * 2]; // 最大孩子数：2t
            this.keys = new int[t * 2 - 1]; // 最大关键字树：2t-1
        }

        @Override
        public String toString() { // 打印所有有效关键字
            return Arrays.toString(Arrays.copyOfRange(keys, 0, keyNumber));
        }

        // 多路查找
        Node get(int key) {
            int i = 0;
            while (i < keyNumber) {
                if (keys[i] == key) { // 找到了
                    return this;
                }
                if (keys[i] > key) {
                    break;
                }
                i++;
            }
            // 循环结束：keys[i] > key 或者 i == keyNumber
            if (leaf) {
                return null; // 如果本节点是叶子节点，没有找到直接返回null
            }
            return children[i].get(key); // 往i个孩子继续查找
        }

        // 往keys的index位置插入key
        void insertKey(int index, int key) {
            // 原index元素后移
            System.arraycopy(keys, index, keys, index + 1, keyNumber - index);
            // 插入key
            keys[index] = key;
            // keyNumber+1
            keyNumber++;
        }

        // 往children的index位置插入child
        void insertChild(int index, Node child) {
            // 原index元素后移
            System.arraycopy(children, index, children, index + 1, keyNumber - index);
            // 插入child
            children[index] = child;
        }

        // 移出指定index处的key
        int removeKey(int index) {
            int removed = keys[index];
            // index后元素往前移动
            // System.arraycopy(keys, index + 1, keys, index, (keyNumber - 1) - index);
            // keyNumber--;
            System.arraycopy(keys, index + 1, keys, index, --keyNumber - index);
            return removed;
        }

        // 移出最左边的key
        int removeLeftmostKey() {
            return removeKey(0);
        }

        // 移出最右边的key
        int removeRightmostKey() {
            return removeKey(keyNumber - 1);
        }

        // 移出index处的child
        Node removeChild(int index) {
            Node removed = children[index];
            System.arraycopy(children, index + 1, children, index, keyNumber - index);
            children[keyNumber] = null; // help gc
            return removed;
        }

        // 移出最左边的child
        Node removeLeftmostChild() {
            return removeChild(0);
        }

        // 移出最右边的child
        Node removeRightmostChild() {
            return removeChild(keyNumber);
        }

        // 获得index孩子左处的兄弟
        Node childLeftSibling(int index) {
            return index > 0 ? children[index - 1] : null;
        }

        // 获得index孩子右处的兄弟
        Node childRightSibling(int index) {
            return index < keyNumber ? children[index + 1] : null;
        }

        // 以追加的方式复制当前节点的所有key和child到target
        void moveToTarget(Node target) {
            int start = target.keyNumber;
            // 复制所有child
            if (leaf) {
                System.arraycopy(children, 0, target.children, start, children.length);
            }
            // 复制所有key
            for (int i = 0; i < keys.length; i++) {
                target.keys[target.keyNumber++] = keys[i];
            }
        }
    }
}
